import torch
import torch.nn as nn
from vortex.network.layers.common import BasicConv, Upsample, SpatialPyramidPooling
from vortex.network.backbone import cspdarknet53


def make_three_conv(in_channels, channels_list):
    m = nn.Sequential(
        BasicConv(in_channels, channels_list[0], 1),
        BasicConv(channels_list[0], channels_list[1], 3),
        BasicConv(channels_list[1], channels_list[0], 1)
    )
    return m


def make_five_conv(in_channels, channels_list):
    m = nn.Sequential(
        BasicConv(in_channels, channels_list[0], 1),
        BasicConv(channels_list[0], channels_list[1], 3),
        BasicConv(channels_list[1], channels_list[0], 1),
        BasicConv(channels_list[0], channels_list[1], 3),
        BasicConv(channels_list[1], channels_list[0], 1)
    )
    return m


def yolo_head(in_channels, channels_list):
    m = nn.Sequential(
        BasicConv(in_channels, channels_list[0], 3),
        BasicConv(channels_list[0], channels_list[1], 1)
    )
    return m


class YoloBody(nn.Module):
    def __init__(self, num_anchors, num_classes):
        super(YoloBody, self).__init__()
        self.backbone = cspdarknet53(None)
        # P3: 256 channel
        # P4: 512 channel
        # P5: 1024 channel
        self.conv_three_1 = make_three_conv(1024, [512, 1024])
        self.spp = SpatialPyramidPooling()
        self.conv_three_2 = make_five_conv(2048, [512, 1024])

        self.upsample1 = Upsample(512, 256)  # P5->P4
        self.conv_for_P4 = BasicConv(512, 256, 1)
        self.conv_five_1 = make_five_conv(512, [256, 512])

        self.upsample2 = Upsample(256, 128)  # P4->P3
        self.conv_for_P3 = BasicConv(256, 128, 1)
        self.conv_five_2 = make_five_conv(256, [128, 256])

        final_out_channels = num_anchors * (5 + num_classes)

        self.yolo_head3 = yolo_head(128, [256, final_out_channels])

        self.downsample1 = BasicConv(128, 256, 3, stride=2)
        self.conv_five_3 = make_five_conv(512, [256, 512])
        self.yolo_head2 = yolo_head(256, [512, final_out_channels])

        self.downsample2 = BasicConv(256, 512, 3, stride=2)
        self.conv_five_4 = make_five_conv(1024, [512, 1024])
        self.yolo_head1 = yolo_head(512, [1024, final_out_channels])

    def forward(self, x):
        x2, x1, x0 = self.backbone(x)
        # x2, x1, x0 for P3, P4, P5
        
        P5 = self.conv_three_1(x0)
        P5 = self.spp(P5)
        P5 = self.conv_three_2(P5)

        P5_upsample = self.upsample1(P5)
        P4 = self.conv_for_P4(x1)
        P4 = torch.cat([P4, P5_upsample], axis=1)
        P4 = self.conv_five_1(P4)

        P4_upsample = self.upsample2(P4)
        P3 = self.conv_for_P3(x2)
        P3 = torch.cat([P3, P4_upsample], axis=1)
        P3 = self.conv_five_2(P3)

        P3_downsample = self.downsample1(P3)
        P4 = torch.cat([P4, P3_downsample], axis=1)
        P4 = self.conv_five_3(P4)

        P4_downsample = self.downsample2(P4)
        P5 = torch.cat([P5, P4_downsample], axis=1)
        P5 = self.conv_five_4(P5)

        out2 = self.yolo_head3(P3)
        out1 = self.yolo_head2(P4)
        out0 = self.yolo_head1(P5)

        return out0, out1, out2
